import heapq
import random

import numpy as np


class Similar:
    def cal_sim(self, vector_a, vector_b):
        """计算两个向量的相似度
        """
        raise NotImplementedError

    def cal_sims(self, vector_a, vectors):
        """计算一个向量与一组向量的相似度，返回一组相似度
        """
        if not isinstance(vector_a, np.ndarray):
            sims = [-999.0] * len(vectors)
            return sims
    
        sims = []
        for a_vector in vectors:
            if not isinstance(a_vector, np.ndarray):
                sim = -999.0
            else:
                sim = self.cal_sim(vector_a, a_vector)
            sims.append(sim)
        return sims
    
    def match_one(self, vector_a, vectors, threshold=0.6, top_k=3):
        """从一组向量中匹配一个与目标向量相近的向量
        Args
        ----
        vector_a : 目标向量
        vectors : 一系列待匹配的向量
        threshold : 相似度阈值
        top_k : 从topo_k 中随机抽取一个
        """
        sims = self.cal_sims(vector_a, vectors)
        # 大于某个阈值
        sims = [(sim, i) for i, sim in enumerate(sims) if sim > threshold]        
        if len(sims) == 0:
            return None, None

        # 选择 top_k 个
        sims = heapq.nlargest(top_k, sims, key=lambda x: x[0])
        # 从中随机抽一个
        sim, ind = random.choice(sims)
        return sim, ind



class CosSimilar(Similar):
    """余弦相似度
    """

    def cal_sim(self, vector_a, vector_b):
        """计算两个向量的相似度
        """
        norm =  (np.linalg.norm(vector_a) * np.linalg.norm(vector_b))
        if norm == 0:
            return -999
        sim = np.dot(vector_a, vector_b) / norm
        return sim

    def cal_sims(self, vector_a, vectors):
        """计算一个向量与一组向量的相似度，返回一组相似度
        Args
        ----
        vector_a : np.ndarray, 单个向量, shape为 (dim, )
        vectors : np.ndarray, 多个向量， shape为 (N, dim)

        Returns
        -------
        sims : np.ndarray, vector_a分别与vectors里每个向量的余弦相似度,
               shape 为（N,）
        """
        assert isinstance(vector_a, np.ndarray), "vetor_a 必须为np.array类型"
        assert isinstance(vectors, np.ndarray), "vetors 必须为np.array类型"
        assert vector_a.shape[0] == vectors.shape[1], "vector_a 和vectors的shape\
                                            关系必须满足(dim,)与(N,dim)关系"

        norm =  (np.linalg.norm(vector_a) * np.linalg.norm(vectors, axis=1))
        sims = np.dot(vector_a, vectors.T) / norm
        return sims


class EuclideanSimilar(Similar):
    def cal_sim(self, vector_a, vector_b):
        sim = np.linalg.norm(vector_b - vector_a)
        return  sim


